import numpy as np
import matplotlib.pyplot as plt

# Parameters
n = 64  # Number of grid points
h = 1 / n  # Step size
omega = 2 / 3  # Weight for weighted Jacobi
tolerance = 1e-2  # Tolerance for convergence (reduce error by a factor of 100)
max_iter = 100  # Maximum number of iterations

# Initialize error reduction data
k_values = np.arange(1, n)  # Wavenumbers from 1 to 63
iterations_jacobi = np.full_like(k_values, max_iter, dtype=int)
iterations_weighted = np.full_like(k_values, max_iter, dtype=int)

# Function to perform Jacobi iteration
def jacobi_iteration(u, f, omega=1.0):
    u_new = np.copy(u)
    for j in range(1, n-1):
        u_new[j] = (1 - omega) * u[j] + (omega / 2) * (u[j-1] + u[j+1] + h**2 * f[j])
    return u_new

x = np.array(np.linspace(0, 1, n))  # Grid points including boundaries

for idx in k_values:
    # Initial guess: mode w_k 
    u = [np.sin(idx * np.pi * x_i) for x_i in x]  # Initial guess

    error = np.linalg.norm(u)

    # Regular Jacobi
    u_jacobi = np.copy(u)
    for it in range(max_iter):
        u_jacobi = jacobi_iteration(u_jacobi, np.zeros_like(u_jacobi), omega=1.0)
        error_jacobi = np.linalg.norm(u_jacobi)
        if error_jacobi < tolerance * error:
            iterations_jacobi[idx] = it + 1
            break

    # Weighted Jacobi
    u_weighted = np.copy(u)
    for it in range(max_iter):
        u_weighted = jacobi_iteration(u_weighted, np.zeros_like(u_weighted), omega=omega)
        error_weighted = np.linalg.norm(u_weighted)
        if error_weighted < tolerance * error:
            iterations_weighted[idx] = it + 1
            break

# Plot results
plt.figure(figsize=(12, 6))

# Regular Jacobi
plt.subplot(1, 2, 1)
plt.plot(k_values, iterations_jacobi, 'bo-', label='Regular Jacobi')
plt.title('Regular Jacobi')
plt.xlabel('Wavenumber $k$')
plt.ylabel('Iterations')
plt.grid(True)
plt.legend()

# Weighted Jacobi
plt.subplot(1, 2, 2)
plt.plot(k_values, iterations_weighted, 'ro-', label=f'Weighted Jacobi ($\\omega = {omega:.2f}$)')
plt.title(f'Weighted Jacobi ($\\omega = {omega:.2f}$)')
plt.xlabel('Wavenumber $k$')
plt.ylabel('Iterations')
plt.grid(True)
plt.legend()

plt.tight_layout()
plt.savefig('9.21.png', dpi=300, bbox_inches='tight')
plt.show()